#!/usr/bin/python3
import logging
import threading
import time

import uservo
import serial
import time

from MotionGroup import MotionGroup
from ServoPose import Pose
import logging

from wrist_gripper import WristGripperControl

from robot_arm import ArmMotionControl

# 标志
exitFlag = 0
recordFlag = 0

# Global const para
THREAD_MODE_CONTROL = ["control", "Control", "CONTROL"]
THREAD_MODE_RECORD = ["record", "Record", "Records", "RECORD", "RECORDS"]

SHORT_TIME_DELAY = 0.02  # time delay para
MEDIA_TIME_DELAY = 0.10
LONG_TIME_DELAY = 1.0

# recording
SAMPLING_FREQUENCY = 10  # units: Hz
fileName = "robot_arm_records.txt"

VEL = {'VERY_LOW_SPEED': 10, 'LOW_SPEED': 15, 'MID_SPEED': 30, 'HIGH_SPEED': 60}
CONTROL_MODE = {'Teach': True, 'Repeat': False}

# DUAL_ARM_INIT_POSE = [140.0, 50.0, 70.0, 0, 2.0, -76.5, 60.0, 0, 0, 0]

# left arm + left wrist
# right arm + right wrist + left and right gripper angle
# DUAL_ARM_START_POSE = [140.0, 50.0, -30.0, 85, 2, -82, -20, 90, 60, 60]
DUAL_ARM_START_POSE = [140.0, 54.0, -30.0, 85,
                       -1, -84, -40, 90, 60, 60]
# 左臂抬升二关节抬升至中间位置
DUAL_ARM_PRO_POSE_LEFT = [140.0, -30.0, -30.0, 85,
                          -1, -84, -40, 90, 60, 60]
# 左臂第一关节向左摆动90度，左臂第二关节俯仰角保持不变
DUAL_ARM_TARGET_POSE_LEFT = [50.0, -30.0, -30.0, 85,
                             -1, -84, -40, 90, 60, 60]

# DUAL_ARM_PRO_POSE_RIGHT = [140.0, 50.0, -30.0, 85,
#                        4, 10, -35, 90, 60, 60]
#
# DUAL_ARM_TARGET_POSE_RIGHT = [140.0, 50.0, -30.0, 85,
#                        80, 10, -35, 90, 60, 60]

# 双臂抬升预定位置
DUAL_ARM_PRO_POSE = [140.0, -30.0, -30.0, 85,
                     -1, 2, -40, 90, 60, 60]

# left and right arm test target pose
DUAL_ARM_TARGET_POSE = [50.0, -30.0, -30.0, 85,
                        89, 2, -40, 90, 60, 60]

ending_servo_idx = [0, 1, 2, 3]

threadLock = threading.Lock()


class dfrThread(threading.Thread):
    def __init__(self, threadID, name, controlPeriod, samplingFrequency, dual_arm_obj):
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.threadName = name
        self.controlPeriod = controlPeriod
        self.samplingFrequency = samplingFrequency
        self.dual_arm_obj = dual_arm_obj

    def run(self):
        print("开始线程：" + self.threadName)
        if (self.threadName in THREAD_MODE_CONTROL):
            logging.info("%s: %s" % (self.threadName, time.ctime(time.time())))

            robot_arm_control(self.threadName, self.dual_arm_obj, self.controlPeriod)

        elif self.threadName in THREAD_MODE_RECORD:
            record_servo_info(self.threadName, self.dual_arm_obj, samplingFrequency=SAMPLING_FREQUENCY)
        else:
            logging.error("self.threadName should be in THREAD_MODE_RECORD or THREAD_MODE_CONTROL !")

        print("退出线程：" + self.name)


def robot_arm_control(threadName, dual_arm, controlPeriod=3000):
    # dual arm motion control threading
    dual_arm_servo = dual_arm.get_arm_servo_list()
    while dual_arm and (threadName in THREAD_MODE_CONTROL):
        keyInput = input('waitKey to continue input (c or x):')
        if keyInput == 'x':
            dual_arm.set_all_servo_mode(CONTROL_MODE="CIRCLE")
            logging.warning('Interrupt:  The waitKey input (c or x): {}'.format(keyInput))
            break
        elif keyInput == 'c':
            # dual_arm再现
            # move the left and right robot arm to start pose for later trajectory control
            for SERVO_ID in dual_arm_servo:
                dual_arm.set_single_servo_pose(
                    Pose=generate_pose_vtc(angle=DUAL_ARM_START_POSE[SERVO_ID], vel=VEL.get('VERY_LOW_SPEED'),
                                           iterval=5000,
                                           t_acc=45,
                                           t_dec=40, power=10000, mean_dps=20, servo_id=SERVO_ID))
            time.sleep(5 * LONG_TIME_DELAY)

            while True:
                loadingKeyInput = input('waitKey to continue input (y or n):')
                if loadingKeyInput == 'n':
                    # 修改标志位，线程上锁
                    # threadLock.acquire()
                    break
                elif loadingKeyInput == 'y':
                    # threadLock.acquire()
                    recordFlag = 1  # record 0->1
                    exitFlag = 0  # exit 0->0
                    # threadLock.release()

                    for i in range(5):

                        # target pose trajectory control for left arm servo id = 0,1,2
                        for SERVO_ID in dual_arm_servo:
                            dual_arm.set_single_servo_pose(
                                Pose=generate_pose_vtc(angle=DUAL_ARM_PRO_POSE_LEFT[SERVO_ID],
                                                       vel=VEL.get('VERY_LOW_SPEED'),
                                                       iterval=controlPeriod,
                                                       t_acc=80,
                                                       t_dec=60, power=10000, mean_dps=20, servo_id=SERVO_ID))
                        time.sleep(3 * LONG_TIME_DELAY)
                        # start pose trajectory control
                        for SERVO_ID in dual_arm_servo:
                            dual_arm.set_single_servo_pose(
                                Pose=generate_pose_vtc(angle=DUAL_ARM_TARGET_POSE_LEFT[SERVO_ID],
                                                       vel=VEL.get('VERY_LOW_SPEED'),
                                                       iterval=controlPeriod,
                                                       t_acc=80,
                                                       t_dec=60, power=10000, mean_dps=20, servo_id=SERVO_ID))
                        time.sleep(3 * LONG_TIME_DELAY)

                    # target pose trajectory control for right arm servo id = 0,1,2,4,5,6
                    time.sleep(10)

                    for i in range(5):
                        # target pose trajectory control for left and right arm
                        for SERVO_ID in dual_arm_servo:
                            dual_arm.set_single_servo_pose(
                                Pose=generate_pose_vtc(angle=DUAL_ARM_PRO_POSE[SERVO_ID], vel=VEL.get('VERY_LOW_SPEED'),
                                                       iterval=controlPeriod,
                                                       t_acc=80,
                                                       t_dec=60, power=10000, mean_dps=20, servo_id=SERVO_ID))
                        time.sleep(3 * LONG_TIME_DELAY)

                        # start pose trajectory control
                        for SERVO_ID in dual_arm_servo:
                            dual_arm.set_single_servo_pose(
                                Pose=generate_pose_vtc(angle=DUAL_ARM_TARGET_POSE[SERVO_ID],
                                                       vel=VEL.get('VERY_LOW_SPEED'),
                                                       iterval=controlPeriod,
                                                       t_acc=50,
                                                       t_dec=40, power=10000, mean_dps=20, servo_id=SERVO_ID))
                        time.sleep(3 * LONG_TIME_DELAY)
        else:
            logging.warning('waitKey input (c or x): {}'.format(keyInput))


# record robot arm servo info
def record_servo_info(threadName, dual_arm, samplingFrequency):
    while dual_arm and samplingFrequency and (threadName in THREAD_MODE_RECORD):
        loadingKeyInput = input('waitKey to record continue input (y or n):')
        if loadingKeyInput == 'n':
            threadName.exit()
            logging.info("exitFlag = True: threadName.exit() !")
            break

        elif loadingKeyInput == 'y':
            with open(fileName, 'w') as f:
                for frame in range(800):
                    # data file write as txt
                    f.write("{}  {}  {}  {} ".format(frame,
                                                            time.ctime(time.time()),
                                                            dual_arm.read_single_servo_angle(SERVO_ID=0) -
                                                            DUAL_ARM_START_POSE[0],
                                                            dual_arm.read_single_servo_angle(SERVO_ID=4) -
                                                            DUAL_ARM_START_POSE[4]) + '\r')
                    time.sleep(1 / samplingFrequency)
            break
        # logging.info("%s: %s" % (threadName, time.ctime(time.time())))
        # recoding the angle data and the power data

# 增加关节舵机角度的vel和iterval参数动态调整接口API
def generate_pose_vtc(angle, servo_id, vel=100, iterval=5000, t_acc=35, t_dec=35, power=10000, mean_dps=70.0):
    pose_tmp = {'angle': angle, 'is_mturn': False, 'interval': iterval,
                'vel': vel, 't_acc': t_acc,
                't_dec': t_dec, 'power': power,
                'mean_dps': mean_dps}
    gen_pose = Pose(POSE=pose_tmp, SERVO_ID=servo_id)
    return gen_pose


def main():
    # main function to construct new control and records threading

    # dual_arm class obj init
    dual_arm = ArmMotionControl(arm_name="dual_arm", dof=6, CONTROL_MODE="DAMPING", with_wrist=False, is_debug=True,
                                damping_power=1000)
    wrist_gripper = WristGripperControl(servo_channel=ending_servo_idx, is_debug=True)

    # 初始化舵机管理器类对象、串口端口地址
    dual_arm.creat_uservo_obj(SERVO_BAUDRATE=115200)
    if dual_arm.servo_buffer(is_scan_servo=True):
        dual_arm.logging_arm_info()
    # dual arm and gripper init
    dual_arm.arm_state_init(force_angle_init=False)
    dual_arm.set_all_servo_mode(CONTROL_MODE="CIRCLE")
    wrist_gripper.all_ending_init(motion_group=4)
    time.sleep(2)

    # 创建新线程
    controlThreadName = "control"
    recordThreadName = "record"

    thread1 = dfrThread(threadID=1,
                        name=controlThreadName,
                        controlPeriod=2000,
                        samplingFrequency=SAMPLING_FREQUENCY,
                        dual_arm_obj=dual_arm)
    thread2 = dfrThread(threadID=2,
                        name=recordThreadName,
                        controlPeriod=None,
                        samplingFrequency=SAMPLING_FREQUENCY,
                        dual_arm_obj=dual_arm)

    # 开启新线程
    thread1.start()
    thread2.start()
    # 等待线程完成
    thread1.join()
    thread2.join()
    logging.info("退出主线程")


# main function
if __name__ == '__main__':
    # 创建一个logger_debug
    import sys

    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)  # Log等级总开关-debug
    print("Logging Set Info: logger.setLevel(logging.DEBUG)  # Log等级总开关-debug")
    try:
        main()
    except KeyboardInterrupt:
        print("Shutting down robot arm debug.")
